# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from coin.base.param_util import to_tuple
from experimental.prophet.graph import graph
from experimental.prophet.ops.null import null
from experimental.prophet.ops.util import to_data_type_proto


def cond(pred, true_value, false_value):
  pred = graph.as_variable(pred).cast(bool)

  if true_value is None:
    assert false_value
    false_value = [graph.as_variable(var) for var in to_tuple(false_value)]
    true_value = [null(var.dtype) for var in false_value]

  elif false_value is None:
    assert true_value
    true_value = [graph.as_variable(var) for var in to_tuple(true_value)]
    false_value = [null(var.dtype) for var in true_value]

  else:
    true_value = [graph.as_variable(var) for var in to_tuple(true_value)]
    false_value = [graph.as_variable(var) for var in to_tuple(false_value)]

  assert len(true_value) == len(false_value)
  for idx in range(len(true_value)):
    assert true_value[idx].dtype == false_value[idx].dtype

  node = graph.create_node('Conditional', [pred, *true_value, *false_value],
                           attrs={'dtype': [to_data_type_proto(var.dtype) for var in true_value]})

  if len(true_value) == 1:
    return node.outputs[0]
  else:
    return node.outputs
